package com.eight.cloud.client;

import com.eight.cloud.client.common.Constant;
import com.eight.cloud.client.config.RequestConfig;
import com.eight.cloud.client.handler.SessionConnectHandler;
import com.eight.cloud.client.service.subscribe.SubscribeObject;
import com.eight.cloud.client.util.RSAUtils;
import org.springframework.messaging.converter.MappingJackson2MessageConverter;
import org.springframework.messaging.converter.MessageConverter;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.web.socket.WebSocketHttpHeaders;
import org.springframework.web.socket.client.WebSocketClient;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;
import org.springframework.web.socket.messaging.WebSocketStompClient;

import java.util.HashMap;
import java.util.Map;

/**
 * @ClassName: WsClient
 * @Author: TXC
 * @Date: 2024-11-19 16:09
 **/

public class WebSocketSubscribeClient {
    private final WebSocketStompClient stompClient;
    private final SessionConnectHandler sessionConnectHandler;
    private final WebSocketHttpHeaders headers;
    private final String serverWsUrl;

    private WebSocketSubscribeClient(RequestConfig requestConfig, SessionConnectHandler sessionConnectHandler,
                                    MessageConverter messageConverter, long[] heartbeat, TaskScheduler taskScheduler) {
        // 创建 WebSocket 客户端
        WebSocketClient webSocketClient = new StandardWebSocketClient();
        // 创建 STOMP 客户端
        WebSocketStompClient stompClient = new WebSocketStompClient(webSocketClient);
        stompClient.setMessageConverter(messageConverter);
        if (taskScheduler != null) {
            stompClient.setDefaultHeartbeat(heartbeat);
            stompClient.setTaskScheduler(taskScheduler);
        }
        this.stompClient = stompClient;
        this.serverWsUrl = requestConfig.getServerWsUrl() + "/websocket/ws";
        this.headers = new WebSocketHttpHeaders();
        headers.add(Constant.SECRET_ID, requestConfig.getSecretId());
        headers.add(Constant.SECRET_KEY, RSAUtils.encrypt(requestConfig.getSecretKey()));
        this.sessionConnectHandler = sessionConnectHandler;
        initializeSessionConnectHandler(this.sessionConnectHandler, this.stompClient, headers, serverWsUrl);
    }

    private void initializeSessionConnectHandler(SessionConnectHandler sessionConnectHandler,
                                                 WebSocketStompClient stompClient, WebSocketHttpHeaders headers,
                                                 String serverWsUrl) {
        sessionConnectHandler.setStompClient(stompClient);
        sessionConnectHandler.setHeaders(headers);
        sessionConnectHandler.setServerWsUrl(serverWsUrl);
    }

    public void connect() {
        stompClient.connect(serverWsUrl, headers, sessionConnectHandler);
    }

    public void disconnect() {
        stompClient.stop();
    }

    public static class Builder {
        private RequestConfig requestConfig;
        private MessageConverter messageConverter;
        private final Map<String, SubscribeObject> subscribeObjects = new HashMap<>();
        private long[] heartbeat;
        private TaskScheduler taskScheduler;

        public Builder setRequestConfig(RequestConfig requestConfig) {
            this.requestConfig = requestConfig;
            return this;
        }

        public Builder setMessageConverter(MessageConverter messageConverter) {
            this.messageConverter = messageConverter;
            return this;
        }

        public Builder addSubscribeObject(SubscribeObject subscribeObject) {
            this.subscribeObjects.put(subscribeObject.getTopic(), subscribeObject);
            return this;
        }

        public Builder setHeartbeat(long[] heartbeat) {
            this.heartbeat = heartbeat;
            return this;
        }

        public Builder setTaskScheduler(TaskScheduler taskScheduler) {
            this.taskScheduler = taskScheduler;
            return this;
        }


        public WebSocketSubscribeClient build() {
            if (messageConverter == null) {
                messageConverter = new MappingJackson2MessageConverter();
            }
            return new WebSocketSubscribeClient(requestConfig, new SessionConnectHandler(subscribeObjects), messageConverter, heartbeat, taskScheduler);
        }
    }
}
